package com.billdimmick.merkabah;

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.URLEncoder;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Date;
import java.util.Iterator;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;

import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.httpclient.HttpClient;
import org.apache.commons.httpclient.HttpConnectionManager;
import org.apache.commons.httpclient.HttpMethod;
import org.apache.commons.httpclient.MultiThreadedHttpConnectionManager;
import org.apache.commons.httpclient.NameValuePair;
import org.apache.commons.httpclient.URI;
import org.apache.commons.httpclient.methods.GetMethod;
import org.apache.commons.httpclient.params.HttpClientParams;

import org.apache.commons.lang.StringUtils;
import org.apache.commons.lang.Validate;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

/**
 * Adapter that manages the requests to SQS, including all of the logic that
 * is needed to sign and encrypt requests.
 * <p> 
 * Consumers of this class may change the runtime configuration on the fly,
 * with the usual expected results.  This allows those consumers to switch their
 * communication from encrypted, signed, plain, or both, depending on their
 * current needs.
 * </p> 
 * @author Bill Dimmick <me@billdimmick.com>
 *
 */
public class SQSManager {

	private static final String SQS_ACTION_KEY = "Action";

	private static final Log log=LogFactory.getLog(SQSManager.class);
	
	private final static String SQS_HOST="queue.amazonaws.com";
	
	private final static String SIGNATURE_METHOD="HmacSHA256";
	private final static String SIGNATURE_VERSION="2";
	private final static String VERSION="2008-01-01";
	
	private final static SimpleDateFormat ISO8601FORMATTER=new SimpleDateFormat("yyyy-MM-dd'T'HH:mm:ssZ"); 
	
	private final SQSSecurityAdapter security;
	private final HttpConnectionManager connectionManager;
	private final HttpClient client;

	/**
	 *  Constructs a new manager.
	 * @param credentials the adapter providing the AWS credentials
	 * @param securityPolicy the security policy for this manager
	 */
	public SQSManager(final SQSSecurityAdapter credentials) {
		Validate.notNull(credentials, "Credential adapter must not be null");
		this.security=credentials;
		Validate.notEmpty(credentials.getAWSAccessId(), "Credential adapter must not supply null access id");
		Validate.notEmpty(credentials.getAWSSecretKey(), "Credential adapter must not supply null secret key");
		
		this.connectionManager=new MultiThreadedHttpConnectionManager();		
		this.client=new HttpClient(this.connectionManager);
		this.client.getParams().setParameter(HttpClientParams.USER_AGENT, "Merkabah/1.0");
	}
		
	public SQSSecurityAdapter getSecurity() {
		return this.security;
	}
	
	/**
	 * Deletes the provided message from the queue.
	 * @param message the message object - must not be null and must not have a null receipt
	 * @throws SQSException if the message cannot be deleted or communications to SQS are exceptional 
	 * @throws IllegalArgumentException if the message is null or the receipt is null
	 */
	public void delete(final Message message) throws SQSException {
		Validate.notNull(message);
		Validate.notNull(message.getReceiptHandle());
		final SortedMap<String, String> parameters=new TreeMap<String, String>();
		parameters.put(SQS_ACTION_KEY, "DeleteMessage");
		parameters.put("ReceiptHandle", message.getReceiptHandle());
		exec(message.getQueue(), parameters);
	}

	/**
	 * Receives a message from the provided queue. 
	 * @return the next available message or null if no message is available
	 * @throws SQSException if the message cannot be retrieved or communications to SQS are exceptional
	 */
	public Message receive(final String queue) throws SQSException {
		final Iterator<Message> results=receive(queue, 1).iterator();
		if (results.hasNext()) {
			return results.next();
		} else {
			return null;
		}
	}

	/**
	 * 
	 * @param max
	 * @return
	 * @throws SQSException if the messages cannot be retrieved or communications to SQS are exceptional
	 */
	public Collection<Message> receive(final String queue, final int max) throws SQSException {		
		return receive(queue, max, -1);
	}

	/**
	 * 
	 * @param max
	 * @param timeout
	 * @return
	 * @throws SQSException if the messages cannot be retrieved or communications to SQS are exceptional
	 */
	public Collection<Message> receive(final String queue, final int max, final int timeout) throws SQSException {
		final SortedMap<String, String> parameters=new TreeMap<String, String>();
		parameters.put(SQS_ACTION_KEY, "ReceiveMessage");		 
		parameters.put("MaxNumberOfMessages", Integer.toString(max));
		if (timeout>0) {
			parameters.put("VisibilityTimeout", Integer.toString(timeout));
		}
		
		final Document document=exec(queue, parameters);
		final Collection<Message> result=new ArrayList<Message>();
		final NodeList list=document.getElementsByTagName("Message");
		for (int i=0; i<list.getLength(); i++) {
			final Element e=(Element)list.item(i);
			final String id=getValue(e, "MessageId");
			final String handle=getValue(e, "ReceiptHandle");
			final String body=getValue(e, "Body");
			result.add(new Message(queue, parseMessage(body), id, handle));
		}
		
		return result;
	}

	/**
	 * Sends a message over SQS, including performing any encryption and signing of the body
	 * @param message the message to send - messages with a null body will send a blank body
	 * @throws SQSException if the message cannot be sent or communications to SQS are exceptional
	 */
	public void send(final Message message) throws SQSException {
		final SortedMap<String, String> parameters=new TreeMap<String, String>();
		parameters.put(SQS_ACTION_KEY, "SendMessage");
		parameters.put("MessageBody", formatMessage(message.getMessage()));
		final Document document=exec(message.getQueue(), parameters);
		final String messageId=getValue(document, "MessageId");
		message.setMessageId(messageId);
	}
	
	/**
	 * Sends a message over SQS, including performing any encryption and signing of the body
	 * @param queue the queue n which to put this message
	 * @param body the message body to send 
	 * @throws SQSException if the message cannot be sent or communications to SQS are exceptional
	 */
	public void send(final String queue, final String body) throws SQSException {
		Validate.notNull(queue, "Queue name cannot be null.");
		Validate.notNull(body, "Provided message body cannot be null.");
		final Message m=new Message(queue, body);
		send(m);
	}
	
	/**
	 * Creates a queue with the specified name and a default visibility timeout of 30s.
	 * @param name the name of the queue - must not be blank or <code>null</code>
	 * @throws SQSException if the queue cannot be created or communications to SQS are exceptional
	 */
	public void createQueue(final String name) throws SQSException {
		createQueue(name, 30);
	}
	
	/**
	 * Creates a queue with the specified name and visibility timeout.
	 * @param name the name of the queue - must not be blank or <code>null</code>
	 * @param timeout the default visibility timeout of the queue - must be a positive integer
	 * @throws SQSException if the queue cannot be created or communications to SQS are exceptional
	 */
	public void createQueue(final String name, final int timeout) throws SQSException {
		Validate.notEmpty(name);
		Validate.isTrue(timeout>0);
		final SortedMap<String, String> parameters=new TreeMap<String, String>();
		parameters.put(SQS_ACTION_KEY, "CreateQueue");
		parameters.put("QueueName", createQueueName(name));
		parameters.put("DefaultVisibilityTimeout", Integer.toString(timeout));
		exec(parameters);
	}

	/**
	 * Deletes a queue.
	 * @param name the name of the queue - must not be blank or <code>null</code>
	 * @throws SQSException if the queue cannot be deleted or communications to SQS are exceptional
	 */
	public void deleteQueue(final String name) throws SQSException {
		Validate.notEmpty(name);
		final SortedMap<String, String> parameters=new TreeMap<String, String>();
		parameters.put(SQS_ACTION_KEY, "DeleteQueue");
		exec(name, parameters);
	}

	/**
	 * Get the approximate number of messages for the given queue.
	 * @param name  the name of the queue - must not be blank or <code>null</code>
	 * @return the approximate number of messages in the queue - zero or more 
	 * @throws SQSException if the communications to SQS are exceptional
	 */
	public int getApproximateNumberOfMessages(final String name) throws SQSException {
		Validate.notEmpty(name);
		final SortedMap<String, String> parameters=new TreeMap<String, String>();
		parameters.put(SQS_ACTION_KEY, "GetQueueAttributes");
		parameters.put("AttributeName", "ApproximateNumberOfMessages");
		final Document document=exec(name, parameters);
		NodeList list=document.getElementsByTagName("Attribute");
		if (list.getLength()>0) {
			Element current=(Element)list.item(0);
			while (current!=null) {
				if ("ApproximateNumberOfMessages".equals(getValue(current, "Name"))) {
					return Integer.parseInt(getValue(current, "Value"));
				}
				current = (Element) current.getNextSibling();
			}
		}
		return 0;
	}

	/**
	 * Lists the queues available.  This method will not return the correct names if
	 * queue name obfuscation is enabled.
	 * @return the list of queues - never <code>null</code>, but may be empty
	 * @throws SQSException if the communications to SQS are exceptional
	 */
	public Collection<String> listQueues() throws SQSException  {
		return listQueues(null);
	}

	/**
	 * Lists the queues available that start with the given prefix.  
	 * This method will not return the correct names if queue name obfuscation is enabled.
	 * @param prefix the prefix to use in the search - <code>null</code> is the same as no prefix
	 * @return the list of queues - never <code>null</code>, but may be empty
	 * @throws SQSException if the communications to SQS are exceptional
	 */
	public Collection<String> listQueues(final String prefix) throws SQSException  {
		final SortedMap<String, String> parameters=new TreeMap<String, String>();
		parameters.put(SQS_ACTION_KEY, "ListQueues");
		if (prefix!=null) {
			parameters.put("QueueNamePrefix", prefix);
		}		
		final Document document=exec(parameters);
		
		final Collection<String> names=getValues(document, "QueueUrl");
		final Collection<String> result=new ArrayList<String>();
		for (String name:names) {
			result.add(name.replaceFirst(String.format("http://%s/", SQS_HOST), ""));
		}
		return result;
	}

	/**
	 * Executes the given request with the provided parameters.
	 * @param parameters the parameters
	 * @return the XML DOM for the response
	 * @throws SQSException if the communications to SQS are exceptional
	 */
	private Document exec(final SortedMap<String, String>parameters) throws SQSException  {
		return exec(null, parameters);
	}
	
	/**
	 * Executes the given request with the provided parameters on the given queue.
	 * @param queue the queue name to use - <code>null</code> indicates that this is a "global" request
	 * @param parameters the parameters to use - must contain "Action" at least
	 * @return the XML DOM for the response
	 * @throws SQSException if the communications to SQS are exceptional
	 */
	private Document exec(final String queue, final SortedMap<String, String>parameters) throws SQSException  {
		Validate.isTrue(parameters.containsKey(SQS_ACTION_KEY));
		try {
			final String realqueue=createQueueName(queue);
			setStandardParameters(parameters);		
			setSignature(realqueue, parameters);
			final HttpMethod method=submit(realqueue, parameters);
			if (method.getStatusCode()!=200) {
				createException(method);
			}
			return createDocument(method);
		} catch (Exception e) {
			if (e instanceof SQSException) {
				throw (SQSException)e;
			} else {
				throw new ClientSideSQSException(e);
			}
		}
	}
	
	/**
	 * Sets the "standard" parameters 
	 * @param parameters the parameter map - must not be <code>null</code>
	 */
	private void setStandardParameters(final Map<String, String> parameters) {
		Validate.notNull(parameters);		
		parameters.put("AWSAccessKeyId", security.getAWSAccessId());
		parameters.put("SignatureVersion", SIGNATURE_VERSION);
		parameters.put("SignatureMethod", SIGNATURE_METHOD);
		parameters.put("Version", VERSION);
		parameters.put("Timestamp", ISO8601FORMATTER.format(new Date()));
	}
	
	/**
	 * Sets the signature for this request. 
	 * @param queue the name of the queue - may be <code>null</code> 
	 * @param parameters the parameters to use when creating this signature
	 * @throws NoSuchAlgorithmException if the HMAC signature algorithm cannot be found - not technically feasible 
	 * @throws InvalidKeyException if the key provided is not usable for the HMAC - not technically feasible
	 * @throws UnsupportedEncodingException if the JVM does not support UTF-8 - not technically feasible
	 */
	protected void setSignature(final String queue, final SortedMap<String, String> parameters) throws NoSuchAlgorithmException, InvalidKeyException, UnsupportedEncodingException {
		final StringBuffer buffer=new StringBuffer("GET\n");
		buffer.append(SQS_HOST);
		buffer.append("\n");
		buffer.append("/");
		buffer.append(StringUtils.defaultString(queue, ""));
		buffer.append("\n");
		buffer.append(join(parameters));
		parameters.put("Signature", this.security.hmac(buffer.toString(), SIGNATURE_METHOD));
	}
	
	/**
	 * Creates one "joined" HTTP parameter string for all of the provided parameters. 
	 * @param parameters the parameters to use - may not be null, but may be empty
	 * @return the joined string - never null, although it may be blank
	 * @throws UnsupportedEncodingException if the JVM does not support UTF-8 - not technically feasible
	 */
	private String join(final Map<String, String> parameters) throws UnsupportedEncodingException {
		final Iterator<Map.Entry<String, String>> values=parameters.entrySet().iterator();
		final StringBuffer buffer=new StringBuffer();
		while (values.hasNext()) {
			final Map.Entry<String, String> entry=values.next();
			buffer.append(entry.getKey());
			buffer.append("=");
			buffer.append(URLEncoder.encode(entry.getValue(), "UTF-8"));
			if (values.hasNext()) {
				buffer.append("&");
			}
		}
		return buffer.toString();
	}
	
	/**
	 * Submits a request to SQS with the provided parameters
	 * @param queue the queue name - may be <code>null</code>
	 * @param parameters the parameters - may be <code>null</code>
	 * @return the resulting response from the request 
	 * @throws IOException if communications to SQS failed  
	 * @throws SAXException if the provided document cannot be parsed
	 * @throws ParserConfigurationException if the document parser cannot be instantiated - not technically feasible
	 * @throws SQSException if the communications to SQS are exceptional   
	 *  
	 */
	protected HttpMethod submit(final String queue, final Map<String, String> parameters) throws IOException, SQSException, ParserConfigurationException, SAXException {
		final HttpMethod method=makeMethod(queue, parameters);
		this.client.executeMethod(method);
		return method;
	}
	
	/**
	 * Creates a new HttpMethod for the provided queue and parameters
	 * @param queue the queue name - may be <code>null</code>
	 * @param parameters the parameters - may be <code>null</code>
	 * @return a new HttpMEthod
	 * @throws IOException if building the URI from the queue has failed
	 */
	protected HttpMethod makeMethod(final String queue, final Map<String, String> parameters) throws IOException{
		final GetMethod method=new GetMethod();
		method.setURI(new URI(String.format("http://%s/%s", SQS_HOST, StringUtils.defaultString(queue, "")), true));
		method.setQueryString(makePairs(parameters));		
		return method;
	}
	
	/**
	 * 
	 * @param parameters
	 * @return
	 */
	private NameValuePair[] makePairs(final Map<String, String> parameters) {
		final NameValuePair pairs[]=new NameValuePair[parameters.size()];
		int i=0;
		for (final Map.Entry<String, String> entry: parameters.entrySet()) {
			pairs[i]=new NameValuePair(entry.getKey(), entry.getValue());
			++i;
		}
		return pairs;
	}
	
	/**
	 * Creates an exception from the provided HTTP output
	 * @param method the HTTP method that generated this exception - must not come from an HTTP 200 OK response
	 * @throws IOException if the URI cannot be retrieved of the document cannot be read
	 * @throws SAXException if the XML document cannot be parsed 
	 * @throws ParserConfigurationException if the XML parser cannot be created - not technically feasible
	 * @throws SQSException if the communications to SQS are exceptional
	 */
	private void createException(final HttpMethod method) throws SQSException, ParserConfigurationException, SAXException, IOException {
		Validate.isTrue(method.getStatusCode()!=200);
		if (log.isWarnEnabled()) {
			log.warn(String.format("Encountered exception for '%s'", method.getURI().toString()));
		}
		
		final Document document=createDocument(method);
		final String type=getValue(document, "Type");
		final String code=getValue(document, "Code");
		final String message=getValue(document, "Message");
		final String id=getValue(document, "RequestId");
		if ("Sender".equals(type)) {
			throw new SenderSQSException(code, message, id);
		} else if ("Receiver".equals(type)) {
			throw new ReceiverSQSException(code, message, id);
		} else throw new IllegalStateException(String.format("Got illegal exception type '%s'", type));
	}
	
	/**
	 * Creates an XML document from an HttpMethod  
	 * @param method the method to use
	 * @return the XML document
	 * @throws ParserConfigurationException if the parse cannot be instantiated - not technically feasible 
	 * @throws IOException if reading the document fails  
	 * @throws SAXException if the document cannot be parsed  
	 */
	private Document createDocument(final HttpMethod method) throws ParserConfigurationException, SAXException, IOException {
		Validate.notNull(method);
		final DocumentBuilder builder=DocumentBuilderFactory.newInstance().newDocumentBuilder();
		final Document result=builder.parse(method.getResponseBodyAsStream());
		return result;
	}
		
	/**
	 * Gets the value of the provided document.
	 * @param document the document to examine - may not be null
	 * @param nodename the name of the node - may not be null
	 * @return the value of the node or <code>null</code> if the node is not found
	 */
	private static String getValue(final Document document, final String nodename) {
		Validate.notNull(document, "The provided document may not be null.");
		return getValue(document.getDocumentElement(), nodename);
	}
	
	/**
	 * Gets a collection of text node values from a document.
	 * @param document the document to examine - may not be null
	 * @param nodename the name of the node - may not be null
	 * @return a collection of values - never null, possibly empty
	 */
	private static Collection<String> getValues(final Document document, final String nodename) {
		Validate.notNull(document, "The provided document may not be null.");
		return getValues(document.getDocumentElement(), nodename);
	}
	
	/**
	 * Gets the value of the provided element.
	 * @param e the element to examine
	 * @param nodename the name of the node - may not be null
	 * @return the value of the node or <code>null</code> if the node is not found
	 */
	private static String getValue(final Element e, final String nodename) {
		final Iterator<String> result=getValues(e, nodename).iterator();
		if (result.hasNext()) {
			return result.next();
		} else {
			return null;
		}		
	}
	
	/**
	 * Gets a collection of text node values from an element.
	 * @param e the element to examine
	 * @param nodename the name of the node - may not be null
	 * @return a collection of values - never null, possibly empty
	 */
	private static Collection<String> getValues(final Element e, final String nodename) {
		Validate.notNull(nodename, "The provided node name may not be null.");
		Validate.notNull(e, "The provided element may not be null.");
		final Collection<String> result=new ArrayList<String>();
		final NodeList list=e.getElementsByTagName(nodename);
		for (int i=0; i<list.getLength(); i++) {
			result.add(list.item(i).getTextContent());
		}
		return result;
	}	
	
	protected String createQueueName(final String original) {
		if (original==null) {
			return null;
		} else if (this.security.shouldObfuscateQueueNames()) {
			return new String(DigestUtils.shaHex(original));
		} else {
			return original;
		}
	}
	
	protected String formatMessage(final String body) throws CryptographicFailureException {
		final MessageFormatAdapter format;
		if (this.security.shouldEncryptMessages()) {
			final String key=this.security.getEncryptionKeyName();
			format=MessageFormatAdapter.create(this.security.encrypt(body, key));
			format.setParameter("EncryptionKey", key);
		} else {
			format=MessageFormatAdapter.create(body);		
		}
		if (this.security.shouldSignMessages()) {
			final String signKey=this.security.getSigningKeyName();
			final String signAlg=this.security.getSigningAlgorithm();
			final String sig=this.security.sign(format.getBody(), signKey, signAlg);
			format.setParameter("Signature", sig);
			format.setParameter("SignAlg", signAlg);
			format.setParameter("SignKey", signKey);
		}
		
		return format.toString();
	}
	
	protected String parseMessage(final String message) throws CryptographicFailureException {
		final MessageFormatAdapter parser=MessageFormatAdapter.parse(message);
		String body=parser.getBody();
		final String key=parser.getParameter("EncryptionKey");
		if (key!=null) {
			body=this.security.decrypt(body, key);
		}
		final String signature=parser.getParameter("Signature");
		if (signature!=null) {
			final String signAlg=parser.getParameter("SignAlg");
			final String signKey=parser.getParameter("SignKey");
			this.security.verify(signature, signKey, signAlg);
		}
		return body;
	}
}
